#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/8/6 20:10
# @Author  : CHEN Wang
# @Site    :
# @File    : index_components.py
# @Software: PyCharm

"""
脚本说明： 获取指数成分股信息, 包含成分股调入调出，权重变动等信息
"""

import pandas as pd
from quant_researcher.quant.project_tool.celebrity import get_lst_condition
from quant_researcher.quant.project_tool.db_operator import db_conn, my_mysql
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.datasource_fetch.common_data_api import t_trade_date

T_INDEX_COMPONENTS_WEIGHT = 'idx_components_weight'
T_STOCK_SECTOR_COMPONENTS = 'astk_qt_sectorcomponents'
T_STOCK_SECTOR_INFO = 'astk_qt_sectorinfo'
T_INDEX_COMPONENTS = 'idx_components'
EARLIEST_START_DATE = '1900-01-01'


def get_concept_hist_stock_pool(concept_code):
    """
    获取概念的成分股信息

    :param str or list concept_code: 概念名称，支持单个(str)、多个(list)
    :return: pd.DataFrame
    """

    conn = db_conn.get_stock_conn()
    sql_str = ('select distinct substr(a.sec_code,3,length(a.sec_code)-2) as code '
               f'from {T_STOCK_SECTOR_COMPONENTS} a,'
               f'{T_STOCK_SECTOR_INFO} b '
               "where b.sector_type = 2 "
               "and a.sector_code = b.sector_code "
               f"and b.sector_name in {tuple(concept_code + [''])}")
    stock_pool_df = my_mysql.read_v3(sql_str, conn)
    conn.close()
    if stock_pool_df.empty:
        LOG.error(f"没有找到概念{concept_code}的信息")
        return
    return stock_pool_df


def _get_cleaned_index_components(index_code):
    """
    获得指数截止到目前所有的成分股信息

    :param str or list index_code: 指数代码，支持单个(str)多个(list)
    :return: pd.DataFrame
    """
    select = ['index_code', 'secu_code as stock_code', 'add_date', 'out_date', 'component_state']
    where = []
    if isinstance(index_code, str):
        where.append(f'index_code=\'{index_code}\'')
    elif isinstance(index_code, list):
        where.append(get_lst_condition(index_code, 'index_code'))
    else:
        LOG.error('目前index_code只支持str或者list，请按要求输入')
        raise NotImplementedError

    conn = db_conn.get_basic_data_conn()
    df = my_mysql.read_v2(select=select, where=where, sfrom=T_INDEX_COMPONENTS, conn=conn)
    conn.close()

    if df.empty:
        return

    return df


# 花花已修复历史数据，不再需要处理脏数据
# def _get_cleaned_index_components(index_code):
#     """
#     同个基金可能同一个add_date有多条进出记录, 针对同个add_date那个需要保留那个具有out_date那个
#     例如add_date=2014-01-01， out_date=null;
#     后来被调出add_date=2014-01-01， out_date=2015-01-01, 则只保留后面这条记录
#
#     :param str index_code: str, 股票池/指数代码 如 '000300'
#     :return:
#     """
#     conn = db_conn.create_a_factor_db_conn()
#     stock_pool_df = pd.read_sql(f"select distinct secu_code as stock_code, add_date, out_date "
#                                 f"from idx_components "
#                                 f"where index_code = '{index_code}' ", conn)
#
#     hist_stocks = set(stock_pool_df.iloc[:, 0])
#     stock_pool_list = []
#     for stock in hist_stocks:
#         temp = stock_pool_df[stock_pool_df['stock_code'] == stock]
#         add_date_num = len(set(temp['add_date']))
#         if len(temp) > 1 and len(temp) > add_date_num:  # 同个基金同一个add_date有多条记录
#             temp.drop_duplicates(subset=['add_date'], inplace=True)
#         stock_pool_list.append(temp)
#     stock_pool_df = pd.concat(stock_pool_list)
#
#     return stock_pool_df


def get_index_latest_stock_pool(index_code):
    """
    获取指数最新的成分股

    :param index_code: 指数代码，支持单个(str)多个(list)
    :return: pd.DataFrame
    """
    stock_pool_df = _get_cleaned_index_components(index_code)
    if stock_pool_df is None:
        LOG.error(f"没有找到{index_code}的成分股信息")
        return

    def filter_components(df):
        mask_1 = (df['component_state'] == 1)
        mask_2 = (df['out_date'].isna())

        return df[mask_1 & mask_2]

    stock_pool_df = stock_pool_df.groupby('index_code').apply(filter_components).reset_index(drop=True)

    return stock_pool_df


def get_index_hist_stock_pool(index_code, start_date=None, end_date=None):
    """
    获取指数在指定时间区间内历史所有的成分股,
    如果start_date和end_date相同，则获取的是历史某天的指数成分股
    如果start_date和end_date皆为None，则获取的是指数历史所有的成分股信息

    :param index_code: 指数代码，支持单个(str)多个(list)
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :return: pd.DataFrame
    """
    stock_pool_df = _get_cleaned_index_components(index_code)
    if stock_pool_df is None:
        LOG.error(f"没有找到{index_code}的成分股信息")
        return

    if start_date is None:  # 未指定起始时间，则默认为'19000101'
        start_date = '19000101'
    start_date = time_tool.format_date_str(start_date, '%Y%m%d')

    if end_date is None:   # 未指定结束时间，则默认为今天
        end_date = time_tool.get_today()
    end_date = time_tool.format_date_str(end_date, '%Y%m%d')

    def filter_components(df):
        mask_1 = (df['add_date'] <= end_date)  # 入选时间不晚于end_date
        mask_2 = (df['out_date'] >= start_date) | (df['out_date'].isna())  # 出池时间不早起起始时间， 或者从未出池
        return df[mask_1 & mask_2]

    stock_pool_df = stock_pool_df.groupby('index_code').apply(filter_components).reset_index(drop=True)

    return stock_pool_df


def get_index_dynamic_stock_pool(index_code, start_date, end_date, **kwargs):
    """
    获取指定时间区间内每天(交易日或自然日)的指数动态成分股

    :param str index_code: str, 股票池/指数代码 如 '000300'
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param kwargs:
        - bool trading_date，是否是交易日，默认：True, 即交易日
    :return:
    """
    trading_date = kwargs.pop('trading_date', True)
    start_date = time_tool.format_date_str(start_date)
    end_date = time_tool.format_date_str(end_date)
    if trading_date:
        date_df = t_trade_date.get_trade_date_by_mkt(start=start_date, end=end_date)
        date_df['trade_date'] = date_df['trade_date'].apply(
            lambda x: time_tool.format_date_str(x, target_fmt='%Y%m%d'))
    else:
        date_df = pd.DataFrame(pd.date_range(start=start_date, end=end_date),
                               columns=['trade_date'])
        date_df['trade_date'] = date_df['trade_date'].apply(
            lambda x: time_tool.format_date_str(x, target_fmt='%Y%m%d'))

    stock_pool_df = get_index_hist_stock_pool(index_code, start_date, end_date)
    if stock_pool_df is None:
        LOG.error(f"没有找到{index_code}在{start_date}到{end_date}之间的成分股信息")
        return

    def filter_components(df):
        hist_stocks = set(df['stock_code'])

        temp_stock_trade_date = []
        for stock in hist_stocks:
            temp = date_df.copy()
            temp['stock_code'] = stock
            temp_stock_trade_date.append(temp)

        temp_stock_trade_date = pd.concat(temp_stock_trade_date, axis=0)
        dynamic_stocks = temp_stock_trade_date.merge(df, on='stock_code')
        max_trade_date = dynamic_stocks.trade_date.max()
        dynamic_stocks['out_date'] = dynamic_stocks['out_date'].replace([None], max_trade_date)
        dynamic_stocks = dynamic_stocks[dynamic_stocks.eval('add_date <= trade_date <= out_date')]
        dynamic_stocks = dynamic_stocks.loc[:, ['index_code', 'stock_code', 'trade_date']]
        dynamic_stocks.drop_duplicates(inplace=True)
        return dynamic_stocks

    stock_pool_df = stock_pool_df.groupby('index_code').apply(filter_components).reset_index(drop=True)

    return stock_pool_df


def get_index_component_weights(index_code, start_date=None, end_date=None, only_latest=True, **kwargs):
    """
    得到指数的成分股及权重
    
    :param str or list or None index_code: 指数代码，支持单个（str）、多个（list）、所有（None）
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :param kwargs:
        - conn: 数据库连接
    :return: pd.DataFrame
            - index_code: 指数代码
            - end_date: 交易日期
            - stock_code: 股票代码
            - stock_weight: 股票权重，不带百分号
    """
    conn = kwargs.pop('conn', None)

    select = ['index_code', 'trade_date as end_date', 'secu_code as stock_code', 'weight/100 as stock_weight']

    where = []

    # index_code过滤条件
    if isinstance(index_code, str) or isinstance(index_code, list):
        if isinstance(index_code, str):
            index_code = [index_code]
        index_where_condition = get_lst_condition(index_code, 'index_code')
        where.append(index_where_condition)
    elif index_code is None:
        index_where_condition = '1 = 1'
    else:
        LOG.error('目前index_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    # 将日期转换成带杠字符串
    if start_date is None:
        start_date = EARLIEST_START_DATE.replace('-', '')
    else:
        start_date = time_tool.format_date_str(start_date, '%Y%m%d')
    if end_date is None:
        end_date = time_tool.get_today(marker='without_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date, '%Y%m%d')

    tmp_conn = db_conn.get_basic_data_conn() if conn is None else conn
    # 只返回最新数据
    if only_latest:
        select = ', '.join(select)
        sql_str = (f"select {select} from {T_INDEX_COMPONENTS_WEIGHT} as a, "
                   f"(select index_code as code, max(trade_date) as enddate "
                   f"from {T_INDEX_COMPONENTS_WEIGHT} "
                   f"where trade_date <= '{end_date}' "
                   f"and {index_where_condition} "
                   f"group by index_code) as b "
                   f"where a.index_code = b.code "
                   f"and a.trade_date = b.enddate ")
        df = my_mysql.read_v3(sql_str, tmp_conn)
    # 返回一段时间数据的where语句
    else:
        where.append(f"trade_date >= '{start_date}'")
        where.append(f"trade_date <= '{end_date}'")
        df = my_mysql.read_v2(select=select, where=where, sfrom=T_INDEX_COMPONENTS_WEIGHT, conn=tmp_conn)

    if conn is None:
        tmp_conn.close()

    if df.empty:
        LOG.error(f"没有找到{index_code}指数的数据，请检查")
        return

    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


"""
之前的API函数，目前都替换成最新的API函数，以下函数已弃用
"""
# def get_index_component_weights(index_list, only_latest=True,
#                                 start_date=None, end_date=None):
#     """
#     得到一个指数列表中每个指数的成份股及权重信息
#
#     :param list index_list: 指数列表
#     :param bool only_latest: 是否只获取最新, 默认为True
#     :param str start_date: 开始时间 "2020-01-01" （开始时间不传则默认为end_date）
#     :param str end_date: 结束时间 "2020-09-01" (结束时间不传则默认为今天)
#     :return: pd.DataFrame
#                 - index_code: 指数代码
#                 - trade_date: 交易日期 str "20200801"
#                 - stock_code: 股票代码
#                 - stock_weight: 股票权重，不带百分号
#     """
#     conn = db_conn.get_basic_data_conn()
#
#     if end_date is None:
#         end_date = time_tool.get_today(marker='with_n_dash')
#     end_date = end_date.replace('-', '')
#     if not only_latest:
#         if start_date is None:
#             start_date = end_date
#             LOG.error(f"获取一段时间内的指数成分股信息，需要指定start_date, 未指定则start_date=end_date")
#         else:
#             start_date = start_date.replace('-', '')
#
#     index_where_condition = get_lst_condition(index_list, 'index_code')
#
#     if only_latest:
#         index_stock_weights = pd.read_sql(f"select a.index_code, a.trade_date, "
#                                           f"a.secu_code as stock_code, a.weight/100 as stock_weight "
#                                           f"from {T_INDEX_COMPONENTS_WEIGHT} as a, "
#                                           f"(select index_code, max(trade_date) as trade_date "
#                                           f"from {T_INDEX_COMPONENTS_WEIGHT} "
#                                           f"where trade_date <= '{end_date}' "
#                                           f"and {index_where_condition} "
#                                           f"group by index_code) as b "
#                                           f"where a.index_code = b.index_code "
#                                           f"and a.trade_date = b.trade_date ", conn)
#     else:
#         index_stock_weights = pd.read_sql(f"select index_code, trade_date, "
#                                           f"secu_code as stock_code, a.weight/100 as stock_weight "
#                                           f"from {T_INDEX_COMPONENTS_WEIGHT} "
#                                           f"where {index_where_condition} "
#                                           f"and trade_date >= '{start_date}' and trade_date <= '{end_date}' ", conn)
#     conn.close()
#
#     if index_stock_weights.empty:
#         LOG.error(f"未找到{index_list}的成份股信息")
#         return
#
#     return index_stock_weights


if __name__ == '__main__':
    start = '2020-01-01'
    end = '2020-05-01'
    index_code = '000300'

    # %% 测试get_cleaned_index_components
    aaa = _get_cleaned_index_components(index_code=index_code)

    # %% 测试get_index_latest_stock_pool
    bbb = get_index_latest_stock_pool(index_code=index_code)

    # %% 测试get_index_hist_stock_pool
    hist_stock = get_index_hist_stock_pool(index_code, start, end)
    hist_stock_1 = get_index_hist_stock_pool(index_code, '2021-05-01', '2021-05-01')
    hist_stock_2 = get_index_hist_stock_pool(index_code)

    # %% 测试get_index_dynamic_stock_pool
    # dynamic_stock = get_index_dynamic_stock_pool(index_code, start, end, trading_date=True)

    # %% 测试get_index_component_weights
    ccc = get_index_component_weights('000300', '2020-01-01', '2020-10-01', True)
